/**
 * Copyright (c) 2012, Brent Mucci
 * All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without modification, 
 * are permitted provided that the following conditions are met:
 * 
 * 1. Redistributions of source code must retain the above copyright notice, 
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright notice, 
 *    this list of conditions and the following disclaimer in the documentation 
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR 
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON 
 * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, 
 * EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 **/

#include "mcl_bitstream_buffer.hpp"

using namespace mcl;

BitstreamBuffer::BitstreamBuffer(std::istream &a_istream) :
	m_istream(a_istream),
	m_bufferBitPos(0), m_bufferBytePos(0), 
	m_primaryBufferSize(0), m_secondaryBufferSize(0), m_isFlipped(false)
{
	// The file cannot be read, return.
	if(!m_istream.good()) return;
	// Determine the size of the file and store it
	m_istream.seekg(0, std::ios::end);
	m_istreamSize = (size_t)m_istream.tellg();
	// Reset the file pointer position to the beginning
	m_istream.seekg(0, std::ios::beg);
	// Assign the primary and secondary buffer pointers to their starting buffers.
	m_pPrimary = m_primaryStreamBuffer;
	m_pSecondary = m_secondaryStreamBuffer;

	// If the number of bytes remaining in the stream is less than the pre-set BUFFER_SIZE 
	// then read in the remaining bytes otherwise use the BUFFER_SIZE.
	m_primaryBufferSize = m_istreamSize - m_istream.tellg() < BUFFER_SIZE ? m_istreamSize : BUFFER_SIZE;
	// The buffer size should be greater than 0 before attempting to read.
	if(m_primaryBufferSize > 0)
	{
		m_istream.read((char *)m_primaryStreamBuffer, m_primaryBufferSize);
	}
	// Load the secondary buffer (if there is any to be loaded).
	m_secondaryBufferSize = m_istreamSize - m_primaryBufferSize < BUFFER_SIZE ? m_istreamSize - m_primaryBufferSize : BUFFER_SIZE;
	if(m_secondaryBufferSize > 0)
	{
		m_istream.read((char *)m_secondaryStreamBuffer, m_secondaryBufferSize);
	}
}

BitstreamBuffer::~BitstreamBuffer()
{
}

unsigned int BitstreamBuffer::read_bits(unsigned int a_bits)
{
	unsigned int bits = peek(a_bits);
	step(a_bits);
	return bits;
}

unsigned int BitstreamBuffer::peek(unsigned int a_bits)
{
	register unsigned int readBits = 0;

	// ((a_bits & 7) ? 1 : 0) = rounding up to an extra byte. numBytes is the ceiling limit of the number of bytes to be read in.
	unsigned int numBytes = (a_bits >> 3) + ((a_bits & 7) ? 1 : 0);

	// CAUTION: The buffer is assumed to be much less than 4 GB, specifically 4 GB - m_primaryBufferSize.
	if(m_bufferBytePos + (a_bits >> 3) > m_primaryBufferSize)
	{
		// The bits being read in overflow into the secondary buffer. Note: The primary buffer still has bits remaining in this condition.
		for(unsigned int i = 0; i < numBytes; i++)
		{
			if(m_bufferBytePos + i >= m_primaryBufferSize)
			{
				// Overrun needs to occur into the secondary buffer
				if((m_bufferBytePos + i - m_primaryBufferSize) >= m_secondaryBufferSize)
				{
					// Overrun cannot occur into the secondary buffer because there aren't enough bytes to read from.
					// The number of bytes that have been successfully read is equal to i.
					numBytes = i;
					// Terminate the loop since future reads will be unsuccessful.
					break;
				}
				else
				{
					// Overrun has occurred into the secondary buffer, read a byte from it.
					readBits = readBits << 8;
					readBits |= m_pSecondary[(m_bufferBytePos + i) - m_primaryBufferSize];
				}
			}
			else
			{
				// Overrun does not yet occur into the secondary buffer, read from the primary buffer.
				readBits = readBits << 8;
				readBits |= m_pPrimary[m_bufferBytePos + i];
			}
		}
	}
	else if(a_bits > 0)
	{
		// All of The bits being read are from the primary buffer.
		readBits = m_pPrimary[m_bufferBytePos];
		for(unsigned int i = 1; i < numBytes; i++)
		{
			readBits = readBits << 8;
			readBits |= m_pPrimary[m_bufferBytePos + i];
		}
	}
	// To be shifted to the right by the number of bits, within a byte, that have not been read in 
	// (complement of the number of bits within a byte that have been read in).
	readBits = readBits >> (((numBytes << 3) - (a_bits & 7)) & 7);
	return readBits;
}

void BitstreamBuffer::step(unsigned int a_bits)
{
	// Increase the bit position within the buffer by a_bits.
	m_bufferBitPos += a_bits;
	// After 8 bits or more have been read in, increment the byte position by the number of whole bytes that have been read.
	m_bufferBytePos += m_bufferBitPos >> 3;
	// Equivalent to m_bufferBitPos %= 8, only the byte remainder is of interest.
	m_bufferBitPos &= 7;
	// If the number of bytes within the primary buffer have been exhausted, load the next buffer.
	if(m_bufferBytePos >= m_primaryBufferSize)
	{
		// This will calculate the byte overflow that has been read in.
		m_bufferBytePos -= m_primaryBufferSize;
		readNextBuffer();
	}
}

bool BitstreamBuffer::readNextBuffer()
{
	// Swap the buffer pointers and read stream data into the secondary buffer.
	if(!m_isFlipped)
	{
		m_isFlipped = true;
		m_pPrimary = m_secondaryStreamBuffer;
		m_pSecondary = m_primaryStreamBuffer;
	}
	else
	{
		m_isFlipped = false;
		m_pPrimary = m_primaryStreamBuffer;
		m_pSecondary = m_secondaryStreamBuffer;
	}
	m_primaryBufferSize = m_secondaryBufferSize;
	// NOTE: *** The following code will execute twice at the very end of the file, place the above code in this method 
	// into its own method since it is functionally different? ***
	// If the number of bytes remaining in the stream is less than the pre-set BUFFER_SIZE then read in the remaining 
	// bytes otherwise use the BUFFER_SIZE.

	m_secondaryBufferSize = m_istreamSize - m_istream.tellg() < BUFFER_SIZE ? m_istreamSize - m_istream.tellg() : BUFFER_SIZE;

	// The buffer size should be greater than 0 before attempting to read.
	if(m_secondaryBufferSize > 0)
	{
		// *** THE PROBLEM IS HERE, REPLACE m_secondaryStreamBuffer with something appropriate!!! Or else it skips!
		m_istream.read((char *)m_pSecondary, m_secondaryBufferSize);
		return true;
	}
	return false;
}

bool BitstreamBuffer::byte_aligned()
{
	// It is byte aligned if the bit position offset (from within a byte) is 0.
	return !m_bufferBitPos;
}

bool BitstreamBuffer::more_data_in_byte_stream()
{
	// If the primary buffer is of size zero then the stream is exhausted, otherwise true.
	return (m_primaryBufferSize!=0);
}

unsigned int BitstreamBuffer::bitsToAlign()
{
	if(byte_aligned()) return 0;
	return (m_bufferBitPos ^ 7) + 1;
}

bool BitstreamBuffer::more_rbsp_data()
{
	// If there is no more data to be read in to a RBSP, return false.
	if(!more_data_in_byte_stream())
	{
		return false;
	}
	else if(byte_aligned())
	{
		// The byte is aligned, therefore the stop bit should be 1 (msb) followed by seven 0 (lsb); return true. 
		// Return false otherise (b/c something is wrong with the bitstream).
		return peek(8) == 128;
	}

	// Byte is not aligned, scan through from right (lsb) to left (msb) until the first 1 is spotted.
	// Condition 1: If the first 1 is spotted before the last bit position in the loop (right to left scan), 
	//              return true (data exists before the stop bit).
	// Condition 2: Return false if the 1 is spotted ONLY in the msb position at the very end of the loop 
	//              (nothing is wrong, there is just nothing left to read before the stop bit).
	// Condition 3: If no 1 is spotted before the loop finishes then return false (something is wrong with the bitstream).
	unsigned int numBitsAlign = bitsToAlign();
	unsigned int bits = peek(numBitsAlign); // Up to 7 bits will be read, fewer may be returned; minimum is 1.
	unsigned int mask = 1;

	// Scan for more at least a 1, starting from the RHS (LSB). Note: This will skip if the numBitsAlign is 0 (aligned).
	for(unsigned int i = 1; i <= numBitsAlign; i++)
	{
		if(bits & mask)
		{
			// Found a 1 in the bitstream
			if(i == numBitsAlign)
			{
				// The 1 was found in the msb position, condition 2 applies.
				return false;
			}
			// The 1 was found before the msb position, condition 1 applies.
			return true;
		}
		bits = bits >> 1;
	}
	
	// Return false otherwise, most likely something is wrong with the bitstream. Condition 3 is the most likely candidate.
	return false;
}

bool BitstreamBuffer::more_rbsp_trailing_data()
{
	return more_data_in_byte_stream();
}

unsigned int BitstreamBuffer::next_bits(unsigned int n)
{
	return peek(n);
}

/**
Additional decoding functionalities to a bitstream buffer.
This includes reading in bit segments of unknown length between 0 and 32 bits in size 
using various coded methods such as exponential golomb coding.
*/

unsigned int BitstreamBuffer::read_bits_ue()
{
	unsigned int expGolombSuffix = 1;
	unsigned char leadingZeroBits,
		numZeros;
	// The maximum number of leading 0 bits, that include the terminating 1, is 32.
	for(numZeros = 0; numZeros < sizeof(unsigned int) << 3; numZeros++)
	{
		leadingZeroBits = read_bits(1);
		if(leadingZeroBits)
		{
			break;
		}
	}
	// If there are no leading 0s then return 0 (should actually throw an exception)
	if(!numZeros) return 0;
	// expGolombSuffix contains a 1 which will be a leading 1 once shifted, this is to form a binary pattern found in unsigned exp golomb decoding.
	expGolombSuffix = expGolombSuffix << numZeros;
	expGolombSuffix |= read_bits(numZeros);

	// Decode the suffix, do this by subtracting 1 from the terminating 1 followed by the suffix (this is a decode pattern).
	return expGolombSuffix - 1;
}

// Code cloning is acceptable to prevent additional stacking.
signed int BitstreamBuffer::read_bits_se()
{
	unsigned int expGolombSuffix = 1, codeNum, elementReturnValue;
	unsigned char leadingZeroBits,
		numZeros;
	// The maximum number of leading 0 bits, that include the terminating 1, is 32.
	for(numZeros = 0; numZeros < sizeof(unsigned int) << 3; numZeros++)
	{
		leadingZeroBits = read_bits(1);
		if(leadingZeroBits)
		{
			break;
		}
	}
	// If there are no leading 0s then return 0 (should actually throw an exception)
	if(!numZeros) return 0;
	// expGolombSuffix contains a 1 which will be a leading 1 once shifted, this is to form a binary pattern found in unsigned exp golomb decoding.
	expGolombSuffix = expGolombSuffix << numZeros;
	expGolombSuffix |= read_bits(numZeros);
	// Decode the suffix, do this by subtracting 1 from the terminating 1 followed by the suffix (this is a decode pattern).
	codeNum = expGolombSuffix - 1;
	// Convert the code num into its partial return value [(codeNum / 2) = codeNum >> 1]
	elementReturnValue = (codeNum >> 1);
	// Apply a ceiling procedure, if necessary
	elementReturnValue += (codeNum & 1);
	// Apply the sign of the number, if applicable: (-1)*exp(codeNum + 1) = (codeNum ^ 1) << ((sizeof(unsigned int) << 3) - 1)
	return elementReturnValue | ((codeNum ^ 1) << ((sizeof(unsigned int) << 3) - 1));
}

void BitstreamBuffer::rbsp_trailing_bits()
{
	/* rbsp_stop_one_bit = */ read_bits(1);
	while(!byte_aligned())
		/* rbsp_alignment_zero_bit = */ read_bits(1);
}